# -*- encoding: utf-8 -*-

# @File    : count_each_annotation.py
# @Time    : 19-7-1
# @Author  : zjh

r"""
Statistic each image's object distribution
"""

import argparse
import csv
from collections import Counter

from .. import common
from .. import annotation_utils


def get_args():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("annotation_dir",
                        help="The root directory of patches")
    parser.add_argument("output_path",
                        help='The convert mode. Default is "UUID2NAME"')
    parser.add_argument("--label_index_map",
                        help='The convert mode. Default is "UUID2NAME"')
    return parser.parse_args()


FLAG = get_args()


def main():
    ann_ps = common.find_file_recursive(FLAG.annotation_dir, ".xml")
    ann_ps = sorted(ann_ps)

    index_label_map = {}
    if FLAG.label_index_map:
        label_index_map = common.load_json(FLAG.label_index_map)
        if isinstance(label_index_map, list):
            index_label_map = {x['id']: x['class'] for x in label_index_map}
        else:
            index_label_map = {v: k for k, v in label_index_map.items()}
    # add additional classes to label index map
    add_classes = set()
    for p in ann_ps:
        ann = annotation_utils.read_annotation(p)
        add_classes.update((x['name'] for x in ann.get("object", [])
                            if x['name'] not in index_label_map.values()))
    add_classes = sorted(add_classes)
    for c in add_classes:
        index_label_map[len(index_label_map)] = c

    num_classes = len(index_label_map)
    head = ['ImageName'] + \
           [index_label_map[i] for i in range(num_classes)]
    data = [head]
    for p in ann_ps:
        ann = annotation_utils.read_annotation(p)
        names = [x['name'] for x in ann.get("object", [])]
        cnt = Counter(names)

        row = [ann['filename']]
        row += [cnt.get(index_label_map[i]) for i in range(num_classes)]
        data.append(row)

    with open(FLAG.output_path, 'w') as f:
        writer = csv.writer(f)
        writer.writerows(data)


if __name__ == '__main__':
    main()
